package com.authine.cloudpivot.web.api.dubbo;

import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.rpc.*;
import com.authine.cloudpivot.engine.api.constants.Constants;
import com.authine.cloudpivot.web.api.config.RunTimeConfig;
import org.slf4j.MDC;

/**
 * @author 1031
 * @since 2019/5/25 15:24
 */
@Activate(group = {com.alibaba.dubbo.common.Constants.PROVIDER, com.alibaba.dubbo.common.Constants.CONSUMER})
public class DubboTraceIdFilter implements Filter {

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {

        if (RunTimeConfig.INSTANCE.isOpenSkyWalkingAgent()) {
            return invoker.invoke(invocation);
        }

        RpcContext rpcContext = RpcContext.getContext();

        // before
        if (rpcContext.isProviderSide()) {
            // get traceId from dubbo consumer，and set traceId to MDC
            String traceId = rpcContext.getAttachment(Constants.TRACE_ID);
            MDC.put(Constants.TRACE_ID, traceId);
        }

        if (rpcContext.isConsumerSide()) {
            // get traceId from MDC, and set traceId to rpcContext
            String traceId = MDC.get(Constants.TRACE_ID);
            rpcContext.setAttachment(Constants.TRACE_ID, traceId);
        }

        Result result = invoker.invoke(invocation);

        // after
        if (rpcContext.isProviderSide()) {
            // clear traceId from MDC
            MDC.remove(Constants.TRACE_ID);
        }

        return result;
    }

}
